package org.renwei.ml.knn;

import com.google.common.collect.*;

import java.util.Collections;
import java.util.List;

/**
 * Created by renwei on 16/8/18.
 */
public class KNN {


    /**
     * kNN算法
     * @param dimension 数据的维度
     * @param dataSet   数据集（包括属性）
     * @param labels    数据集对应的标签
     * @param k         k值
     * @param inX       待测数据
     * @return
     */
    public static String kNN(int dimension, double[][] dataSet, int[] labels, int k,
        double[] inX) {
        List<LableDis> lLists = Lists.newArrayList();
        for (int i = 0; i < dataSet.length; i++) {
            lLists.add(new LableDis(dimension, dataSet[i], labels[i]));
        }

        //求矩阵距离
        for (int i = 0; i < lLists.size(); i++) {
            LableDis lab = lLists.get(i);
            double sumN = 0;
            for (int j = 0; j < dimension; j++) {
                sumN += Math.pow(lab.getPoints()[j] - inX[j], lab.getPoints().length);
            }
            lab.setDistance(Math.pow(sumN, 1.0 / lab.getPoints().length));
        }

        //按照距离排序
        Collections.sort(lLists);

        //取前k个label中概率最多的label
        Multiset<String> labelMultiset = HashMultiset.create();
        for (int i = 0; i < k; i++) {
            labelMultiset.add(String.valueOf(lLists.get(i).getLabel()));
        }
        Iterable<Multiset.Entry<String>> entriesSortedByCount =
            Multisets.copyHighestCountFirst(labelMultiset).entrySet();

//        for (Multiset.Entry<String> entry : entriesSortedByCount) {
        //            System.out.println(entry.getElement() + " count: " + entry.getCount());
        //        }

        return entriesSortedByCount.iterator().next().getElement();
    }

    public static void main(String[] args) {
        //input
        int dimension = 2;
        double[][] dataSet = {{1, 1.1}, {1, 1}, {0, 0}, {0, 0.1}};
        int[] labels = {1, 1, 2, 2};
        int k = 3;
        double[] inX = {10, 10};
        System.out.println(kNN(dimension, dataSet, labels, k, inX));
    }
}


class LableDis implements Comparable<LableDis> {
    private int dimension;          //属性维度
    private double[] points;        //属性值
    private int label;              //标签
    private double distance;        //距离

    public LableDis(int dimension, double[] points, int label) {
        this.dimension = dimension;
        this.points = points;
        this.label = label;
    }

    public int getDimension() {
        return dimension;
    }

    public void setDimension(int dimension) {
        this.dimension = dimension;
    }

    public double[] getPoints() {
        return points;
    }

    public void setPoints(double[] points) {
        this.points = points;
    }

    public int getLabel() {
        return label;
    }

    public void setLabel(int label) {
        this.label = label;
    }

    public double getDistance() {
        return distance;
    }

    public void setDistance(double distance) {
        this.distance = distance;
    }

    public int compareTo(LableDis o) {
        return ComparisonChain.start().compare(this.distance, o.distance).result();
    }
}
